from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os, random , sys, time

import numpy as np
import tensorflow as tf

from model import Seq2SeqModel
from properties import Properties
from data_proc import PAD_ID, GO_ID


def _assert_lengths(encoder_size, decoder_size, encoder_inputs, decoder_inputs, decoder_masks):
    if len(encoder_inputs) != encoder_size:
        raise ValueError("Encoder length must be equal to the one in bucket,"
                        " %d != %d." % (len(encoder_inputs), encoder_size))
    if len(decoder_inputs) != decoder_size:
        raise ValueError("Decoder length must be equal to the one in bucket,"
                       " %d != %d." % (len(decoder_inputs), decoder_size))
    if len(decoder_masks) != decoder_size:
        raise ValueError("Weights length must be equal to the one in bucket,"
                       " %d != %d." % (len(decoder_masks), decoder_size))


def run_step(sess, props, model, encoder_inputs, decoder_inputs, decoder_masks, bucket_id, fw_only):
    """
    单步运行
    :param sess: session
    :param props: Properties
    :param model: Seq2seq model
    :param encoder_inputs: inputs
    :param decoder_inputs: inputs
    :param decoder_masks: decoder_weights
    :param bucket_id: 哪个bucket
    :param fw_only: 是否不backward
    :return: tuple, Gradient Norm, loss, eval outputs
    """
    buckets = eval(props.getProperties('Attrs.buckets'))
    encoder_size, decoder_size = buckets[bucket_id]
    _assert_lengths(encoder_size, decoder_size,encoder_inputs,decoder_inputs, decoder_masks)
    input_feed = {}
    for step in range(encoder_size):
        input_feed[model.encoder_inputs[step].name] = encoder_inputs[step]
    for step in range(decoder_size):
        input_feed[model.decoder_inputs[step].name] = decoder_inputs[step]
        input_feed[model.decoder_masks[step].name] = decoder_masks[step]

    last_target = model.decoder_inputs[decoder_size].name
    input_feed[last_target] = np.zeros([model.batch_size], dtype=np.int32)

    if not fw_only:
        output_feed = [model.train_ops[bucket_id],
                       model.gradient_norms[bucket_id],
                       model.losses[bucket_id]]
    else:
        output_feed = [model.losses[bucket_id]]
        for step in range(decoder_size):
            output_feed.append(model.outputs[bucket_id][step])

    outputs = sess.run(output_feed, input_feed)
    if not fw_only:
        return outputs[1], outputs[2], None  # Gradient norm, loss, no outputs.
    else:
        return None, outputs[0], outputs[1:]  # No gradient norm, loss, outputs.


def get_batch(props, data, bucket_id, batch_size = None):
    """
    获取一个batch的数据
    :param props: Properties, 参数对象
    :param data: list, 数据
    :param bucket_id: int, 选择的bucket
    :return: tuple
    """
    buckets = eval(props.getProperties('Attrs.buckets'))
    if batch_size == None:
        batch_size = int(props.getProperties('Attrs.batch_size'))
    encoder_size, decoder_size = buckets[bucket_id]
    encoder_inputs, decoder_inputs = [], []
    # Get a random batch of encoder and decoder inputs from data,
    # pad them if needed, reverse encoder inputs and add GO to decoder.
    for _ in range(batch_size):
      encoder_input, decoder_input = random.choice(data[bucket_id])

      # Encoder inputs are padded and then reversed.
      encoder_pad = [PAD_ID] * (encoder_size - len(encoder_input))
      encoder_inputs.append(list(reversed(encoder_input + encoder_pad)))

      # Decoder inputs get an extra "GO" symbol, and are padded then.
      decoder_pad_size = decoder_size - len(decoder_input) - 1
      decoder_inputs.append([GO_ID] + decoder_input +
                            [PAD_ID] * decoder_pad_size)

    # Now we create batch-major vectors from the data selected above.
    batch_encoder_inputs, batch_decoder_inputs, batch_weights = [], [], []

    # Batch encoder inputs are just re-indexed encoder_inputs.
    for length_idx in range(encoder_size):
      batch_encoder_inputs.append(
          np.array([encoder_inputs[batch_idx][length_idx]
                    for batch_idx in range(batch_size)], dtype=np.int32))

    # Batch decoder inputs are re-indexed decoder_inputs, we create weights.
    for length_idx in range(decoder_size):
      batch_decoder_inputs.append(
          np.array([decoder_inputs[batch_idx][length_idx]
                    for batch_idx in range(batch_size)], dtype=np.int32))

      # Create target_weights to be 0 for targets that are padding.
      batch_weight = np.ones(batch_size, dtype=np.float32)
      for batch_idx in range(batch_size):
        # We set weight to 0 if the corresponding target is a PAD symbol.
        # The corresponding target is decoder_input shifted by 1 forward.
        if length_idx < decoder_size - 1:
          target = decoder_inputs[batch_idx][length_idx + 1]
        if length_idx == decoder_size - 1 or target == PAD_ID:
          batch_weight[batch_idx] = 0.0
      batch_weights.append(batch_weight)
    return batch_encoder_inputs, batch_decoder_inputs, batch_weights


def create_cell(type ):
    """
    基本的设置单个cell
    :param type: cell-type
    :return: 正常的 cell
    """
    if type == 'GRU':
        return tf.contrib.rnn.GRUCell
    elif type == 'LSTM':
        return tf.contrib.rnn.BasicLSTMCell

